/te-ci L1 L2 jax

Greptile Summary

Fixed a bug where TransformerEngine/JAX couldn't query mesh information from meshes set via jax.set_mesh(). Previously, the code only checked _PXLA_THREAD_RESOURCES.env.physical_mesh, which is populated when using with mesh: context managers but not when using jax.set_mesh().

Key Changes:

• Added new _get_mesh() helper function that checks both thread resources and jax.sharding.get_abstract_mesh()
• Replaced 6 direct accesses to _PXLA_THREAD_RESOURCES.env.physical_mesh with calls to _get_mesh()
• Ensures backward compatibility by checking thread resources first, then falling back to abstract mesh

Impact:

• Enables TE/JAX to work correctly with both mesh initialization patterns
• No breaking changes to existing code using with mesh: context

Confidence Score: 4/5

• This PR is safe to merge with minimal risk
• The fix is well-structured with a clear fallback pattern. The change is backwards compatible (checking thread resources first maintains existing behavior), localized to one file, and addresses a specific bug. Confidence reduced by 1 point due to one minor consideration around null safety.
• No files require special attention

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant User
    participant JAX
    participant TE as TransformerEngine
    participant PXLA as PXLA Thread Resources
    participant GetAbstractMesh as jax.sharding.get_abstract_mesh()

    Note over User,GetAbstractMesh: Scenario 1: Mesh set via 'with mesh:' context
    User->>JAX: with mesh:
    JAX->>PXLA: Set physical_mesh
    User->>TE: Call TE function (e.g., is_mesh_available())
    TE->>TE: _get_mesh()
    TE->>PXLA: Check physical_mesh
    PXLA-->>TE: Returns mesh (not empty)
    TE-->>User: Returns mesh
    
    Note over User,GetAbstractMesh: Scenario 2: Mesh set via jax.set_mesh()
    User->>JAX: jax.set_mesh(mesh)
    JAX->>GetAbstractMesh: Store mesh in abstract context
    Note over PXLA: physical_mesh remains None or empty
    User->>TE: Call TE function (e.g., is_mesh_available())
    TE->>TE: _get_mesh()
    TE->>PXLA: Check physical_mesh
    PXLA-->>TE: Returns None or empty mesh
    TE->>GetAbstractMesh: Fallback to get_abstract_mesh()
    GetAbstractMesh-->>TE: Returns mesh from jax.set_mesh()
    TE-->>User: Returns mesh